#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
诊断和修复 ChromaDB 路径不一致问题
"""

import os
import sys
import asyncio
import shutil
from pathlib import Path

# 添加项目根目录到Python路径
sys.path.insert(0, os.getcwd())

from src.config import load_config

async def diagnose_chromadb_paths():
    """诊断 ChromaDB 路径问题"""
    print("🔍 诊断 ChromaDB 路径配置问题...")
    
    # 1. 检查配置文件中的路径
    try:
        config = load_config()
        chromadb_config = config.get('databases', {}).get('chromadb', {})
        config_path = chromadb_config.get('persist_directory', './chroma_db')
        print(f"📋 配置文件中的路径: {config_path}")
    except Exception as e:
        print(f"❌ 读取配置失败: {e}")
        return
    
    # 2. 检查实际存在的目录
    base_dir = Path("./data")
    possible_paths = [
        "./chroma_db",
        "./data/chroma_db", 
        "./data/chromadb"
    ]
    
    existing_dirs = []
    for path in possible_paths:
        full_path = Path(path)
        if full_path.exists():
            existing_dirs.append({
                'path': path,
                'full_path': full_path,
                'size': sum(f.stat().st_size for f in full_path.rglob('*') if f.is_file()),
                'file_count': len(list(full_path.rglob('*')))
            })
    
    print(f"\n📁 发现的 ChromaDB 目录:")
    for dir_info in existing_dirs:
        print(f"  - {dir_info['path']}:")
        print(f"    文件数: {dir_info['file_count']}")
        print(f"    大小: {dir_info['size']} bytes")
    
    # 3. 检查每个目录中的集合
    print(f"\n🔍 检查各目录中的 ChromaDB 集合:")
    
    try:
        import chromadb
        from chromadb.config import Settings
        
        for dir_info in existing_dirs:
            print(f"\n  检查目录: {dir_info['path']}")
            try:
                client = chromadb.PersistentClient(
                    path=str(dir_info['full_path']),
                    settings=Settings(anonymized_telemetry=False)
                )
                
                collections = client.list_collections()
                print(f"    集合数量: {len(collections)}")
                
                total_docs = 0
                for collection in collections:
                    count = collection.count()
                    total_docs += count
                    print(f"      - {collection.name}: {count} 个文档")
                
                print(f"    总文档数: {total_docs}")
                
            except Exception as e:
                print(f"    ❌ 检查失败: {e}")
    
    except ImportError:
        print("⚠️  ChromaDB 未安装，跳过集合检查")
    
    # 4. 分析问题并给出建议
    print(f"\n📊 分析结果:")
    
    if len(existing_dirs) > 1:
        print(f"⚠️  发现多个 ChromaDB 目录，可能导致数据分散")
        
        # 找到数据最多的目录
        main_dir = max(existing_dirs, key=lambda x: x['file_count'])
        print(f"📈 数据最多的目录: {main_dir['path']} ({main_dir['file_count']} 个文件)")
        
        # 检查配置路径是否存在
        config_dir_exists = any(d['path'] == config_path for d in existing_dirs)
        if not config_dir_exists:
            print(f"❌ 配置文件指定的路径 {config_path} 不存在或无数据")
            print(f"💡 建议: 将数据迁移到配置路径或更新配置")
        else:
            print(f"✅ 配置路径存在数据")
    
    elif len(existing_dirs) == 1:
        single_dir = existing_dirs[0]
        if single_dir['path'] != config_path:
            print(f"⚠️  唯一的数据目录 {single_dir['path']} 与配置路径 {config_path} 不匹配")
            print(f"💡 建议: 更新配置文件或迁移数据")
        else:
            print(f"✅ 路径配置正确")
    
    else:
        print(f"❌ 未找到任何 ChromaDB 数据目录")
        print(f"💡 建议: 创建配置路径 {config_path}")
    
    return existing_dirs, config_path

async def fix_chromadb_paths():
    """修复 ChromaDB 路径问题"""
    print(f"\n🔧 开始修复 ChromaDB 路径问题...")
    
    existing_dirs, config_path = await diagnose_chromadb_paths()
    
    if not existing_dirs:
        print(f"❌ 没有发现现有数据，无需修复")
        return
    
    # 找到数据最多的目录作为主目录
    main_dir = max(existing_dirs, key=lambda x: x['file_count'])
    main_path = main_dir['full_path']
    
    print(f"\n📋 修复计划:")
    print(f"  主数据目录: {main_dir['path']}")
    print(f"  目标配置路径: {config_path}")
    
    # 如果主目录就是配置路径，只需要合并其他目录
    if str(main_path) == config_path or main_dir['path'] == config_path:
        print(f"✅ 主数据目录已经是配置路径")
        
        # 合并其他目录的数据
        for dir_info in existing_dirs:
            if dir_info['path'] != main_dir['path']:
                print(f"🔄 合并目录 {dir_info['path']} 到主目录...")
                await merge_chromadb_data(dir_info['full_path'], main_path)
    
    else:
        # 需要将主目录迁移到配置路径
        config_full_path = Path(config_path)
        
        print(f"🔄 迁移主数据目录到配置路径...")
        
        # 确保目标目录的父目录存在
        config_full_path.parent.mkdir(parents=True, exist_ok=True)
        
        if config_full_path.exists():
            print(f"⚠️  目标路径已存在，将合并数据")
            await merge_chromadb_data(main_path, config_full_path)
        else:
            print(f"📁 移动 {main_path} 到 {config_full_path}")
            shutil.move(str(main_path), str(config_full_path))
        
        # 合并其他目录
        for dir_info in existing_dirs:
            if dir_info['path'] != main_dir['path']:
                print(f"🔄 合并目录 {dir_info['path']} 到配置路径...")
                await merge_chromadb_data(dir_info['full_path'], config_full_path)
    
    print(f"\n✅ ChromaDB 路径修复完成！")
    print(f"📋 所有数据现在都在: {config_path}")

async def merge_chromadb_data(source_path: Path, target_path: Path):
    """合并 ChromaDB 数据"""
    try:
        import chromadb
        from chromadb.config import Settings
        
        print(f"  🔄 合并 {source_path} -> {target_path}")
        
        # 打开源数据库
        source_client = chromadb.PersistentClient(
            path=str(source_path),
            settings=Settings(anonymized_telemetry=False)
        )
        
        # 打开目标数据库
        target_client = chromadb.PersistentClient(
            path=str(target_path),
            settings=Settings(anonymized_telemetry=False)
        )
        
        # 获取源数据库的所有集合
        source_collections = source_client.list_collections()
        
        for source_collection in source_collections:
            collection_name = source_collection.name
            print(f"    📋 处理集合: {collection_name}")
            
            # 获取源集合的所有数据
            all_data = source_collection.get()
            
            if not all_data['documents']:
                print(f"      ⚠️  集合 {collection_name} 为空，跳过")
                continue
            
            # 获取或创建目标集合
            target_collection = target_client.get_or_create_collection(
                name=collection_name,
                metadata=source_collection.metadata or {}
            )
            
            # 批量添加数据
            try:
                target_collection.add(
                    documents=all_data['documents'],
                    metadatas=all_data['metadatas'],
                    ids=all_data['ids'],
                    embeddings=all_data['embeddings'] if all_data['embeddings'] else None
                )
                print(f"      ✅ 成功合并 {len(all_data['documents'])} 个文档")
            except Exception as e:
                print(f"      ❌ 合并失败: {e}")
                # 尝试逐个添加（处理ID冲突）
                success_count = 0
                for i, (doc, meta, doc_id) in enumerate(zip(
                    all_data['documents'], 
                    all_data['metadatas'], 
                    all_data['ids']
                )):
                    try:
                        embedding = all_data['embeddings'][i] if all_data['embeddings'] else None
                        target_collection.add(
                            documents=[doc],
                            metadatas=[meta],
                            ids=[doc_id],
                            embeddings=[embedding] if embedding else None
                        )
                        success_count += 1
                    except Exception:
                        # ID 冲突或其他错误，跳过
                        pass
                print(f"      ✅ 成功合并 {success_count}/{len(all_data['documents'])} 个文档")
        
        # 删除源目录（如果合并成功）
        if source_path != target_path:
            print(f"    🗑️  删除源目录: {source_path}")
            shutil.rmtree(source_path)
        
    except Exception as e:
        print(f"    ❌ 合并失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(fix_chromadb_paths())